import os

import cv2
import numpy as np
import pytorch_lightning as pl
import torch
import torch.optim as optim

from downstream.criterion import DownstreamLoss
from utils.metrics import compute_IoU_from_cmatrix, confusion_matrix


class LightningTrainImg(pl.LightningModule):
    def __init__(self, model_images, model_classifier, config):
        super().__init__()
        self.model_images = model_images
        self.model_classifier = model_classifier

        self._config = config
        self.ignore_index = config["ignore_index"]
        self.ignore_indexes = config["ignore_indexes"]

        if config["loss"].lower() == "lovasz":
            self.criterion = DownstreamLoss(
                ignore_index=config["ignore_index"],
                device=self.device,
            )
        else:
            self.criterion = torch.nn.CrossEntropyLoss(
                ignore_index=config["ignore_index"],
            )

    def configure_optimizers(self):
        if self._config.get("lr_backbone", None) is not None:
            optimizer = optim.SGD(
                [
                    {
                        "params": iter(self.model_images.encoder.parameters()),
                        "lr": self._config["lr_backbone"],
                    },
                    {"params": iter(self.model_images.decoder.parameters())},
                    {"params": iter(self.model_classifier.parameters())},
                ],
                lr=self._config["lr"],
                momentum=self._config["sgd_momentum"],
                dampening=self._config["sgd_dampening"],
                weight_decay=self._config["weight_decay"],
            )
        else:
            optimizer = optim.SGD(
                [
                    {"params": iter(self.model_images.parameters())},
                    {"params": iter(self.model_classifier.parameters())},
                ],
                lr=self._config["lr"],
                momentum=self._config["sgd_momentum"],
                dampening=self._config["sgd_dampening"],
                weight_decay=self._config["weight_decay"],
            )
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, self._config["num_epochs"]
        )
        return [optimizer], [scheduler]

    def training_step(self, batch, batch_idx):
        img_feats = self.model_images(batch["ims"])
        img_feats = img_feats.permute(0, 2, 3, 1)
        img_probs = self.model_classifier(img_feats)

        loss = self.loss(batch, img_probs)
        self.log(
            "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
        )

        return loss

    def loss(self, batch, img_probs):
        tmp_pred = img_probs[0].argmax(2).cpu()
        tmp_pred = np.array(tmp_pred, dtype=np.uint8) * 21
        # tmp_pred = cv2.applyColorMap(tmp_pred, cv2.COLORMAP_RAINBOW)
        cv2.imwrite("pred.png", tmp_pred)

        tmp_gt = batch["msks"][0].cpu()
        tmp_gt = np.array(tmp_gt, dtype=np.uint8) * 21
        # tmp_gt = cv2.applyColorMap(tmp_gt, cv2.COLORMAP_RAINBOW)
        cv2.imwrite("gt.png", tmp_gt)

        tmp_input = batch["ims"][0].cpu()
        tmp_input = tmp_input.permute(1, 2, 0) * 255
        tmp_input = np.array(tmp_input, dtype=np.uint8)
        cv2.imwrite("input.png", tmp_input)

        gt_msk = batch["msks"]
        return self.criterion(img_probs, gt_msk)

    def validation_step(self, batch, batch_idx):
        img_feats = self.model_images(batch["ims"])
        img_feats = img_feats.permute(0, 2, 3, 1)
        img_probs = self.model_classifier(img_feats)

        loss = self.loss(batch, img_probs)
        self.log(
            "val_loss",
            loss,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            sync_dist=True,
        )

        img_probs = img_probs.flatten(start_dim=0, end_dim=2).argmax(1)
        gt_msk = torch.flatten(batch["msks"])
        c_matrix = confusion_matrix(img_probs, gt_msk, self._config["classes"])
        return loss, c_matrix

    def validation_epoch_end(self, outputs) -> None:
        c_matrix = sum([o[1] for o in outputs])

        m_IoU, fw_IoU, per_class_IoU = compute_IoU_from_cmatrix(
            c_matrix, self.ignore_indexes
        )

        self.log("m_IoU", m_IoU, prog_bar=True, logger=True, sync_dist=True)
        self.log("fw_IoU", fw_IoU, prog_bar=True, logger=True, sync_dist=True)

        return super().validation_epoch_end(outputs)
